"""数据处理"""


def filter_850_vwnd():
    """滤波例子"""

    vwndData = ERA5DataForTimeFilter("vwnd", engine='netcdf4')

    for data, time in vwndData:
        data = cal.reverse_longitude(data)
        data = data.sel(level=850)

        # 3-8天时间滤波
        data = cal.butter_band_pass(data, 3, 8, 4, fs=4)

        values = np.concatenate([data.values[:, :, -80:], 
                                    data.values, 
                                    data.values[:, :, :80]], axis=-1)


        # 改为用傅里叶变换滤波，只要向西的波
        fft = np.fft.fft2(values, axes=(0, 2))
        fft = np.fft.fftshift(fft, axes=(0, 2))
        freqLon = np.fft.fftfreq(values.shape[2], d=CONFIG.dataRes)
        freqLon = np.fft.fftshift(freqLon)
        freqTime = np.fft.fftfreq(values.shape[0])
        freqTime = np.fft.fftshift(freqTime)
        # 找到 -1/18 < freq < 0 的位置
        arg0 = np.searchsorted(freqLon, -CONFIG.lowPassFrequency) 
        arg1 = np.searchsorted(freqLon, 0)
        arg2 = np.searchsorted(freqLon, CONFIG.lowPassFrequency)
        argCenter = np.searchsorted(freqTime, 0)

        fft[:argCenter+1, :, :arg0] = 0
        fft[:argCenter+1, :, arg1:] = 0
        fft[argCenter:, :, :arg1] = 0
        fft[argCenter:, :, arg2:] = 0

        values = np.fft.ifft2(
            np.fft.ifftshift(fft, axes=(0, 2)), axes=(0, 2))\
                .real[:, :, 80:-80]

        data.values = values

        data = data.sel(time=time)

        savePath = MRGTCDataManager.filter_850_vwnd_folder / \
            f"FLTv.850.{time[0:4]}{time[-2:]}.nc"

        data.to_netcdf(savePath, encoding=int16_encoding(data))

            